from datetime import datetime
from typing import List, Optional, Dict, Any
from sqlalchemy.orm import Session
from sqlalchemy import desc, and_

from app.models.chat import Conversation, Message, ConversationSummary, MessageRole, ChatMode
from app.core.database import get_db
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, BaseMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI
import os


class ChatMemoryService:
    """聊天记忆服务，负责数据库操作和摘要生成"""
    
    def __init__(self):
        self.llm = ChatOpenAI(
            model="deepseek-chat",
            temperature=0.7,
            max_tokens=2048,
            api_key=os.getenv("OPENAI_API_KEY"),
            base_url=os.getenv("OPENAI_BASE_URL"),
        )
        # 摘要触发阈值（消息数量）
        self.summary_threshold = 20
        # 保留的最近消息数量（不被摘要）
        self.keep_recent_messages = 10
    
    def create_conversation(self, db: Session, title: Optional[str] = None, user_id: Optional[int] = None, chat_mode: ChatMode = ChatMode.NORMAL) -> Conversation:
        """创建新对话"""
        conversation = Conversation(
            title=title or f"对话 {datetime.now().strftime('%Y-%m-%d %H:%M')}",
            user_id=user_id,
            chat_mode=chat_mode
        )
        db.add(conversation)
        db.commit()
        db.refresh(conversation)
        return conversation
    
    def get_conversation(self, db: Session, conversation_id: int) -> Optional[Conversation]:
        """获取对话"""
        return db.query(Conversation).filter(Conversation.id == conversation_id).first()
    
    def get_latest_conversation_by_mode(self, db: Session, user_id: int, chat_mode: ChatMode) -> Optional[Conversation]:
        """获取用户指定模式的最新对话"""
        return db.query(Conversation).filter(
            and_(Conversation.user_id == user_id, Conversation.chat_mode == chat_mode)
        ).order_by(desc(Conversation.updated_at)).first()
    
    def save_message(self, db: Session, conversation_id: int, role: MessageRole, content: str, metadata: Optional[Dict[str, Any]] = None) -> Message:
        """保存消息到数据库"""
        # 获取当前对话的消息数量作为序号
        message_count = db.query(Message).filter(Message.conversation_id == conversation_id).count()
        
        message = Message(
            conversation_id=conversation_id,
            role=role,
            content=content,
            sequence_number=message_count + 1,
            metadata=str(metadata) if metadata else None
        )
        
        db.add(message)
        
        # 更新对话的消息计数和最后活跃时间
        conversation = db.query(Conversation).filter(Conversation.id == conversation_id).first()
        if conversation:
            conversation.message_count = message_count + 1
            conversation.last_activity_at = datetime.now()
        
        db.commit()
        db.refresh(message)
        return message
    
    def get_conversation_messages(self, db: Session, conversation_id: int, limit: Optional[int] = None) -> List[Message]:
        """获取对话的所有消息"""
        query = db.query(Message).filter(Message.conversation_id == conversation_id).order_by(Message.sequence_number)
        if limit:
            query = query.limit(limit)
        return query.all()
    
    def get_recent_messages(self, db: Session, conversation_id: int, count: int = 10) -> List[Message]:
        """获取最近的消息"""
        return db.query(Message).filter(
            Message.conversation_id == conversation_id
        ).order_by(desc(Message.sequence_number)).limit(count).all()[::-1]  # 反转以保持时间顺序
    
    async def generate_summary(self, messages: List[Message]) -> str:
        """生成消息摘要"""
        if not messages:
            return ""
        
        # 构建消息文本
        message_text = "\n".join([
            f"{msg.role.value}: {msg.content}" for msg in messages
        ])
        
        # 摘要提示词
        prompt = ChatPromptTemplate.from_template(
            """请对以下对话内容进行简洁的摘要，保留关键信息和上下文：

{messages}

摘要："""
        )
        
        parser = StrOutputParser()
        chain = prompt | self.llm | parser
        
        summary = await chain.ainvoke({"messages": message_text})
        return summary
    
    async def save_summary(self, db: Session, conversation_id: int, messages: List[Message]) -> ConversationSummary:
        """生成并保存摘要"""
        summary_content = await self.generate_summary(messages)
        
        summary = ConversationSummary(
            conversation_id=conversation_id,
            summary_content=summary_content,
            start_message_seq=messages[0].sequence_number,
            end_message_seq=messages[-1].sequence_number,
            summary_type="partial"
        )
        
        db.add(summary)
        db.commit()
        db.refresh(summary)
        return summary
    
    def get_active_summaries(self, db: Session, conversation_id: int) -> List[ConversationSummary]:
        """获取有效的摘要"""
        return db.query(ConversationSummary).filter(
            and_(
                ConversationSummary.conversation_id == conversation_id,
                ConversationSummary.is_active == True
            )
        ).order_by(ConversationSummary.start_message_seq).all()
    
    async def should_create_summary(self, db: Session, conversation_id: int) -> bool:
        """判断是否需要创建摘要"""
        message_count = db.query(Message).filter(Message.conversation_id == conversation_id).count()
        return message_count >= self.summary_threshold
    
    async def get_conversation_context(self, db: Session, conversation_id: int) -> List[BaseMessage]:
        """获取对话上下文（摘要 + 最近消息）"""
        context_messages = []
        
        # 获取有效摘要
        summaries = self.get_active_summaries(db, conversation_id)
        if summaries:
            # 将摘要作为系统消息添加
            summary_text = "\n\n".join([s.summary_content for s in summaries])
            context_messages.append(SystemMessage(content=f"对话历史摘要：\n{summary_text}"))
        
        # 获取最近的消息
        recent_messages = self.get_recent_messages(db, conversation_id, self.keep_recent_messages)
        
        # 转换为LangChain消息格式
        for msg in recent_messages:
            if msg.role == MessageRole.USER:
                context_messages.append(HumanMessage(content=msg.content))
            elif msg.role == MessageRole.ASSISTANT:
                context_messages.append(AIMessage(content=msg.content))
            elif msg.role == MessageRole.SYSTEM:
                context_messages.append(SystemMessage(content=msg.content))
        
        return context_messages
    
    async def process_and_summarize(self, db: Session, conversation_id: int):
        """处理消息并在需要时创建摘要"""
        if await self.should_create_summary(db, conversation_id):
            # 获取需要摘要的消息（除了最近的消息）
            all_messages = self.get_conversation_messages(db, conversation_id)
            
            # 找到最后一个摘要的结束位置
            last_summary = db.query(ConversationSummary).filter(
                ConversationSummary.conversation_id == conversation_id
            ).order_by(desc(ConversationSummary.end_message_seq)).first()
            
            start_seq = 1
            if last_summary:
                start_seq = last_summary.end_message_seq + 1
            
            # 计算需要摘要的消息范围
            total_messages = len(all_messages)
            end_seq = total_messages - self.keep_recent_messages
            
            if end_seq > start_seq:
                # 获取需要摘要的消息
                messages_to_summarize = [
                    msg for msg in all_messages 
                    if start_seq <= msg.sequence_number <= end_seq
                ]
                
                if messages_to_summarize:
                    await self.save_summary(db, conversation_id, messages_to_summarize)